from sentence_transformers import CrossEncoder


class CrossEncoderReranker:
    def __init__(self, model_name="cross-encoder/ms-marco-MiniLM-L-6-v2"):
        self.model = CrossEncoder(model_name)

    def rerank(self, query, documents, top_n=5):
        """使用交叉编码器对检索结果进行重排序"""
        if not documents:
            return []

        # 准备查询-文档对
        pairs = [(query, doc) for doc in documents]

        # 获取相关性分数
        scores = self.model.predict(pairs)

        # 按分数排序
        ranked = sorted(zip(documents, scores), key=lambda x: x[1], reverse=True)

        # 返回top_n结果
        return [doc for doc, score in ranked[:top_n]]


class RRFReranker:
    def __init__(self, k=60):
        """
        RRF重排序器
        :param k: 平滑常数，默认60（标准RRF值）
        """
        self.k = k

    def rerank(self, vector_results, bm25_results, top_n=5):
        """
        使用RRF算法融合两个检索系统的结果

        :param vector_results: 向量检索结果列表（已排序，[doc1, doc2, ...]）
        :param bm25_results: BM25检索结果列表（已排序，[doc1, doc2, ...]）
        :param top_n: 返回前N个结果
        :return: 重排序后的文档列表
        """
        # 创建文档->排名的映射
        vector_ranks = {doc: rank for rank, doc in enumerate(vector_results, 1)}
        bm25_ranks = {doc: rank for rank, doc in enumerate(bm25_results, 1)}

        # 获取所有唯一文档
        all_docs = set(vector_results) | set(bm25_results)

        # 计算每个文档的RRF分数
        doc_scores = []
        for doc in all_docs:
            # 获取文档在各自系统中的排名（如果不存在则给最大排名）
            rank_vector = vector_ranks.get(doc, len(vector_results) + 1)
            rank_bm25 = bm25_ranks.get(doc, len(bm25_results) + 1)

            # 计算RRF分数
            score = (1 / (rank_vector + self.k)) + (1 / (rank_bm25 + self.k))
            doc_scores.append((doc, score))

        # 按RRF分数降序排序
        doc_scores.sort(key=lambda x: x[1], reverse=True)

        # 返回top_n结果
        return [doc for doc, _ in doc_scores[:top_n]]


# 修改后的retrieve方法
